#include <fstream>
#include <random>
#include <thread>

#include <bthread/bthread.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

#include "util/radix_tree/var_len_data_accessor.h"
#include "util/pool/pool.h"

using namespace hawking::indexlib;
DEFINE_int32(radixtree_test_tasks, 128, "num of threads or bthreads");
std::uniform_int_distribution<int> dis(0, 10000000);
std::default_random_engine e{std::random_device{}()};

void* Write(void* arg) {
    VarLenDataAccessor* var_len_accessor = reinterpret_cast<VarLenDataAccessor*>(arg);
    while (1) {
        std::string str = std::to_string(dis(e));
        std::string_view sv(str.data(), str.length());
        var_len_accessor->AppendValue(sv);
        bthread_usleep(1 * 1000);
    }
}

void* Read(void* arg) {
    VarLenDataAccessor* var_len_accessor = reinterpret_cast<VarLenDataAccessor*>(arg);
    while (1) {
        uint64_t doc_count = var_len_accessor->GetDocCount();
        for (uint32_t idx = 0; idx < doc_count; ++idx) {
            uint8_t* data = nullptr;
            uint32_t length = 0;
            EXPECT_EQ(true, var_len_accessor->ReadData(idx, std::ref(data)));
            // LOG(INFO) << "read doc " << idx << ", str " << std::string((char*)data, length);
            bthread_usleep(1 * 1000);
        }
    }
}

TEST(radix_tree_test, radix_tree_read) {
    std::unique_ptr<MMapAllocator> allocator = std::make_unique<MMapAllocator>();
    std::unique_ptr<Pool> pool = std::make_unique<Pool>(allocator.get(), 10 * 1024 * 1024);
    VarLenDataAccessor var_len_accessor;
    var_len_accessor.Init(pool.get());

    // write bthread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, Write, &var_len_accessor);

    // read bhtreads
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_radixtree_test_tasks);
    for (size_t idx = 0; idx < FLAGS_radixtree_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, Read, &var_len_accessor);
    }

    std::thread stat_thread([&var_len_accessor](){
        uint64_t idx = 0;
        while (1) {
            LOG(INFO) << "idx " << idx++ << ", doc count " << var_len_accessor.GetDocCount();
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }
    });
    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }
    stat_thread.join();
}

int main(int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}
